#include "./spi_flash.h"
#include "../device/gpio.h"
#include "../rtos/rtos.h"
#include <string.h>
#include "./cmd_uart.h"

namespace driver::spi_flash {

using namespace device;
// SPI3
using SCK = PB<3>;
using MISO = PB<4>;
using MOSI = PB<5>;
using CS = PB<6>;

static constexpr uint8_t ReadStatus = 0x05;
static constexpr uint8_t WriteEnable = 0x06;
static constexpr uint8_t ReadId = 0x90;
static constexpr uint8_t SectorErase = 0x20;
static constexpr uint8_t ReadData = 0x03;
static constexpr uint8_t PageProgram = 0x02;

static constexpr uint32_t PageSize = 256;
static constexpr uint32_t SectorSize = 4096;

static uint8_t tmp_buf[SectorSize];

void init() {
    RCC->AHB1ENR |= RCC_AHB1ENR_GPIOBEN;
    RCC->APB1ENR |= RCC_APB1ENR_SPI3EN;

    ConfigAf<6, SCK, MISO, MOSI>();
    ConfigOutput<CS>();

    SPI3->CR1 = 0x00;
    SPI3->CR2 = 0x00;
    SPI3->I2SCFGR = 0x00;
    SPI3->CR1 = 0x0344;
    CS::Set();

    cmdUart->throw_log(0, "flash id:%04x", get_id());
}

static uint8_t spi_write_read(uint8_t val) {
    while (0x00 == (SPI3->SR & SPI_SR_TXE));
    SPI3->DR = val;
    while (0x00 == (SPI3->SR & SPI_SR_RXNE));
    return SPI3->DR;
}

uint16_t get_id() {
    uint16_t id = 0;
    CS::Clr();
    spi_write_read(ReadId);
    spi_write_read(0x00);
    spi_write_read(0x00);
    spi_write_read(0x00);
    id = spi_write_read(0x00);
    id <<= 8;
    id |= spi_write_read(0x00);
    CS::Set();
    return id;
}

void read(uint32_t addr, void *buf, uint32_t len) {
    uint8_t *buf8 = reinterpret_cast<uint8_t *>(buf);
    
    CS::Clr();
    spi_write_read(ReadData);
    spi_write_read(addr >> 16);
    spi_write_read(addr >> 8);
    spi_write_read(addr);
    while (len --) {
        *buf8 = spi_write_read(0xFF);
        buf8 ++;
    }
    CS::Set();
}

static inline uint32_t sector_base_addr(uint32_t addr) {
    return addr & (~0xFFF);
}

static inline uint32_t sector_max_size(uint32_t addr) {
    return SectorSize - (addr & 0xFFF);
}

static inline uint32_t page_base_addr(uint32_t addr) {
    return addr & (~0xFF);
}

static inline uint32_t page_max_size(uint32_t addr) {
    return PageSize - (addr & 0xFF);
}

static void write_enable() {
    CS::Clr();
    spi_write_read(WriteEnable);
    CS::Set();
}

static uint8_t query_status() {
    uint8_t ret ;
    CS::Clr();
    spi_write_read(ReadStatus);
    ret = spi_write_read(0xFF);
    CS::Set();
    return ret;
}

static void erase(uint32_t addr) {
    write_enable();
    CS::Clr();
    spi_write_read(SectorErase);
    spi_write_read(addr >> 16);
    spi_write_read(addr >> 8);
    spi_write_read(addr);
    CS::Set();

    while ( ( query_status() & 0x01 ) != 0 ) {
        osDelay(5);
    }
}

static void page_program(uint32_t addr, const void *buf, uint32_t len) {
    const uint8_t *buf8 = reinterpret_cast<const uint8_t*>(buf);
    write_enable();
    CS::Clr();
    spi_write_read(PageProgram);
    spi_write_read(addr >> 16);
    spi_write_read(addr >> 8);
    spi_write_read(addr);
    while (len --) {
        spi_write_read(*buf8);
        buf8 ++;
    }
    CS::Set();

    while ( (query_status() & 0x01) != 0 ) {
        osDelay(5);
    }
}

static void sector_program(uint32_t addr, const void *buf, uint32_t len) {
    const uint8_t *buf8 = reinterpret_cast<const uint8_t*>(buf);
    for (uint32_t index = 0; index < len; ) {
        uint32_t n = len - index;
        uint32_t max = page_max_size(addr + index);
        if (n > max) {
            n = max;
        }
        page_program(addr + index, buf8 + index, n);
        index += n;
    }
}

static void write_all(uint32_t addr, const void *buf, uint32_t len) {
    const uint8_t *buf8 = reinterpret_cast<const uint8_t*>(buf);
    for (uint32_t index = 0; index < len; ) {
        uint32_t n = len - index;
        uint32_t max = sector_max_size(addr + index);
        if (n > max) {
            n = max;
        }
        sector_program(addr + index, buf8 + index, n);
        index += n;
    }
}

void write(uint32_t addr, const void *buf, uint32_t len) {
    const uint8_t *buf8 = reinterpret_cast<const uint8_t*>(buf);

    uint32_t index = 0;

    for (; index < len;) {
        uint32_t base_addr = sector_base_addr(addr + index);
        uint32_t offset = addr + index - base_addr;
        uint32_t max = sector_max_size(addr + index);
        uint32_t n = len - index;
        if (n > max) {
            n = max;
        }
        read(base_addr, tmp_buf, SectorSize);
        memcpy(tmp_buf + offset, buf8 + index, n);
        erase(base_addr);
        write_all(base_addr, tmp_buf, SectorSize);
        index += n;
    }
}

}

